{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# WatsonX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call `complete` with a prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import WatsonX\n",
    "\n",
    "credentials = {\n",
    "    \"url\": \"https://enter.your-ibm.url\",\n",
    "    \"apikey\": \"insert_your_api_key\",\n",
    "}\n",
    "\n",
    "project_id = \"insert_your_project_id\"\n",
    "\n",
    "resp = WatsonX(credentials=credentials, project_id=project_id).complete(\n",
    "    \"Paul Graham is\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " a programmer, entrepreneur, and author. He is the co-founder of Y Combinator, a startup accelerator. He is also the author of the books Hackers & Painters and On Lisp.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call `chat` with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import ChatMessage, WatsonX\n",
    "\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "\n",
    "resp = WatsonX(\n",
    "    model_id=\"meta-llama/llama-2-70b-chat\",\n",
    "    credentials=credentials,\n",
    "    project_id=project_id,\n",
    ").chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant:  Ahoy matey! Yer lookin' fer a tale of the high seas, eh? Well, settle yerself down with a pint o' grog and listen close, for I've got a story that'll make yer timbers shiver!\n",
      "\n",
      "It were a dark and stormy night, and me ship, the Black Dragon, were sailin' through treacherous waters. The winds were howlin' like a pack o' wolves and the waves were crashin' against the hull like angry giants. Me and me crew were battle-hardened pirates, but even we were gettin' nervous.\n",
      "\n",
      "Suddenly, a great shadow loomed ahead, and we found ourselves face to face with the most fearsome sea monster I'd ever laid eyes on! It were a giant squid, with tentacles as long as me ship and a beak that could snap a man in two.\n",
      "\n",
      "Now, I'd never been one to shy away from a fight, but this beast were different. It were a creature of the deep, and it seemed to have a mind of its own. It attacked us with a fury, wrappin' its tentacles around the ship and pullin' us down into the depths.\n",
      "\n",
      "We fought with all our might, but it were no use. The squid were too powerful, and it were draggin' us down to a watery grave. I knew that we had to come up with a plan, and fast, or we'd be seein' Davy Jones' locker for ourselves!\n",
      "\n",
      "That's when I remembered a trick I'd learned from a wise old sea dog back in port. I told me crew to pour vinegar on the squid's tentacles, and sure enough, it worked like a charm! The squid let out a deafening screech and released its grip on the ship. We were able to sail away, battered but alive.\n",
      "\n",
      "And that, me hearty, is the tale of how I, Blackbeak the pirate, defeated the sea monster and saved me crew. Now, if ye'll excuse me, I've got a hankerin' for some grog and a good sea shanty!\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Using `stream_complete` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import WatsonX\n",
    "\n",
    "llm = WatsonX(credentials=credentials, project_id=project_id)\n",
    "\n",
    "resp = llm.stream_complete(\"Paul Graham is\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " a programmer, entrepreneur, and writer. He is the co-founder of Y Combinator, a startup accelerator. He is also the author of the essay collection Hackers & Painters."
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `stream_chat` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import WatsonX\n",
    "\n",
    "llm = WatsonX(\n",
    "    model_id=\"meta-llama/llama-2-70b-chat\",\n",
    "    credentials=credentials,\n",
    "    project_id=project_id,\n",
    ")\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Arrrr, me hearty! Listen close and I'll spin ye a yarn of the high seas.\n",
      "\n",
      "Ahoy, me name be Captain Blackbeak, the most feared pirate on the seven seas. Me and me crew, the Blackheart Gang, have been sailin' these waters for nigh on 20 years, plunderin' the riches of the landlubbers and bringin' terror to them that dare cross us.\n",
      "\n",
      "One dark and stormy night, we came upon a grand ship, the \"Golden Dolphin,\" sailin' the seas with a cargo hold full of gold, silver, and precious jewels. The captain, a foolish man named Captain Redbeard, thought to outrun us, but me and me crew were hot on his tail.\n",
      "\n",
      "We battled through the ragin' storm, cannons blazin', swords clashin', and after a fierce fight, we claimed the \"Golden Dolphin\" as our own. The loot was divided among the crew, and we set sail for the nearest port to celebrate our victory.\n",
      "\n",
      "But little did we know, Captain Redbeard had a secret weapon, a stowaway on board, a beautiful maiden named Sarah. She was a fierce warrior, skilled in the ways of combat, and she had a plan to take back her ship and its treasure.\n",
      "\n",
      "As we docked at the port, Sarah struck, takin' out me crew one by one, and makin' her way to the treasure room. But I, Captain Blackbeak, was ready for her. I challenged her to a duel, and after a fierce battle, I emerged victorious.\n",
      "\n",
      "Sarah was impressed by me skill and bravery, and she decided to join me crew. Together, we sailed the seas, plunderin' and pillagin' for years to come. And that, me hearty, be the tale of how I, Captain Blackbeak, defeated the fierce warrior Sarah and claimed the \"Golden Dolphin\" as me own.\n",
      "\n",
      "Now, if ye be wantin' to hear another tale of the high seas, just let me know, and I'll be happy to spin ye another yarn."
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import WatsonX\n",
    "\n",
    "llm = WatsonX(\n",
    "    model_id=\"meta-llama/llama-2-70b-chat\",\n",
    "    credentials=credentials,\n",
    "    project_id=project_id,\n",
    "    temperature=0,\n",
    "    max_new_tokens=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = llm.complete(\"Paul Graham is\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " a well-known entrepreneur, investor, and writer. He is best known for his work in the field of startup accelerators, and for his essays on startup success and entrepreneurship.\n",
      "\n",
      "Graham has a background in computer science and philosophy, and has worked in various roles in the tech industry, including as a programmer, software designer, and product manager. He is perhaps best known, however, for his work as a co-founder of Y Com\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-1zyBTJDh-py3.9",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
